import numpy as np
from typing import List, Tuple, Optional


class SwitchPuzzleSolver:
    def __init__(self, grid: List[List[int]]):
        """
        初始化开关谜题求解器

        Args:
            grid: 任意大小的二维数组
            - 0: 无开关
            - 1,2,3: 可手动操作的开关（1=关闭, 2=等待, 3=开启）
            - -1,-2,-3: 不可手动操作的开关（-1=关闭, -2=等待, -3=开启）
        """
        self.grid = np.array(grid)
        self.rows = len(grid)
        self.cols = len(grid[0]) if grid else 0
        self.mod = 3  # 模3运算

    def get_adjacent_positions(self, row: int, col: int) -> List[Tuple[int, int]]:
        """
        获取指定位置的十字相邻位置

        Args:
            row: 行索引
            col: 列索引

        Returns:
            相邻位置列表
        """
        adjacent = []
        directions = [(-1, 0), (1, 0), (0, -1), (0, 1)]  # 上下左右

        for dr, dc in directions:
            new_row, new_col = row + dr, col + dc
            if 0 <= new_row < self.rows and 0 <= new_col < self.cols:
                adjacent.append((new_row, new_col))

        return adjacent

    def is_manual_switch(self, row: int, col: int) -> bool:
        """判断指定位置的开关是否可以手动操作"""
        return self.grid[row, col] > 0

    def get_switch_state(self, row: int, col: int) -> int:
        """获取开关的实际状态（去掉符号）"""
        return abs(self.grid[row, col])

    def create_equation_matrix(self) -> Tuple[np.ndarray, np.ndarray]:
        """
        创建线性方程组的系数矩阵和常数向量

        Returns:
            系数矩阵A和常数向量b
        """
        # 只考虑可手动操作的开关
        manual_switches = []
        for row in range(self.rows):
            for col in range(self.cols):
                if self.is_manual_switch(row, col):
                    manual_switches.append((row, col))

        n_manual = len(manual_switches)
        if n_manual == 0:
            return np.array([]), np.array([])

        A = np.zeros((n_manual, n_manual), dtype=int)
        b = np.zeros(n_manual, dtype=int)

        # 为每个可手动操作的开关创建方程
        for i, (row, col) in enumerate(manual_switches):
            # 目标状态：从当前状态到开启状态(3)需要的改变次数
            current_state = self.get_switch_state(row, col)
            target_state = 3

            # 计算需要的改变次数
            if current_state == 1:  # 关闭→开启：需要2次按动
                required_change = 2
            elif current_state == 2:  # 等待→开启：需要1次按动
                required_change = 1
            elif current_state == 3:  # 已经是开启：需要0次按动
                required_change = 0
            else:
                required_change = 0

            b[i] = required_change

            # 考虑所有开关（包括不可手动操作的）对当前开关的影响
            for j, (switch_row, switch_col) in enumerate(manual_switches):
                # 检查这个开关是否会影响当前开关
                if self.affects_switch(switch_row, switch_col, row, col):
                    A[i, j] = 1

        return A, b

    def affects_switch(
        self, source_row: int, source_col: int, target_row: int, target_col: int
    ) -> bool:
        """判断source位置的开关是否会影响target位置的开关"""
        # 获取source位置的相邻位置（包括自己）
        adjacent_positions = self.get_adjacent_positions(source_row, source_col)
        adjacent_positions.append((source_row, source_col))

        # 检查target是否在影响范围内
        return (target_row, target_col) in adjacent_positions

    def gaussian_elimination_mod3(
        self, A: np.ndarray, b: np.ndarray
    ) -> Optional[np.ndarray]:
        """
        在模3域中进行高斯消元（因为每次按动改变1次状态）

        Args:
            A: 系数矩阵
            b: 常数向量

        Returns:
            解向量，如果无解则返回None
        """
        if len(A) == 0:
            return np.array([])

        n = len(A)
        augmented = np.column_stack([A, b])

        # 前向消元
        for i in range(n):
            # 寻找主元
            pivot_row = i
            for j in range(i + 1, n):
                if abs(augmented[j, i]) > abs(augmented[pivot_row, i]):
                    pivot_row = j

            if pivot_row != i:
                augmented[i], augmented[pivot_row] = (
                    augmented[pivot_row].copy(),
                    augmented[i].copy(),
                )

            # 如果主元为0，尝试找到非零主元
            if augmented[i, i] == 0:
                for j in range(i + 1, n):
                    if augmented[j, i] != 0:
                        augmented[i], augmented[j] = (
                            augmented[j].copy(),
                            augmented[i].copy(),
                        )
                        break
                else:
                    # 检查是否有解
                    if augmented[i, -1] != 0:
                        return None  # 无解
                    continue

            # 消元
            for j in range(i + 1, n):
                if augmented[j, i] != 0:
                    pivot_val = int(augmented[i, i])
                    current_val = int(augmented[j, i])
                    factor = (current_val * pow(pivot_val, -1, 3)) % 3
                    augmented[j] = (augmented[j] - factor * augmented[i]) % 3

        # 回代求解
        x = np.zeros(n, dtype=int)
        for i in range(n - 1, -1, -1):
            if augmented[i, i] == 0:
                if augmented[i, -1] != 0:
                    return None  # 无解
                x[i] = 0  # 自由变量设为0
                continue

            sum_val = 0
            for j in range(i + 1, n):
                sum_val = (sum_val + augmented[i, j] * x[j]) % 3

            pivot_val = int(augmented[i, i])
            const_val = int(augmented[i, -1])
            x[i] = ((const_val - sum_val) * pow(pivot_val, -1, 3)) % 3

        return x

    def solve(self) -> Optional[List[Tuple[int, int]]]:
        """
        求解开关谜题

        Returns:
            需要按动的开关位置列表，如果无解则返回None
        """
        A, b = self.create_equation_matrix()
        solution = self.gaussian_elimination_mod3(A, b)

        if solution is None:
            return None

        if len(solution) == 0:
            return []

        # 将解转换为开关位置
        manual_switches = []
        for row in range(self.rows):
            for col in range(self.cols):
                if self.is_manual_switch(row, col):
                    manual_switches.append((row, col))

        moves = []
        for i, count in enumerate(solution):
            if count > 0:  # 需要按动
                row, col = manual_switches[i]
                for _ in range(count):
                    moves.append((row, col))

        return moves

    def simulate_moves(self, moves: List[Tuple[int, int]]) -> np.ndarray:
        """
        模拟执行移动后的最终状态

        Args:
            moves: 移动列表

        Returns:
            最终状态网格
        """
        final_grid = self.grid.copy()

        for row, col in moves:
            if not self.is_manual_switch(row, col):
                continue

            # 改变当前位置：正负分别循环加一
            v = final_grid[row, col]
            if v > 0:
                final_grid[row, col] = v % 3 + 1
            elif v < 0:
                final_grid[row, col] = -((-v) % 3 + 1)

            # 改变相邻位置
            adjacent_positions = self.get_adjacent_positions(row, col)
            for adj_row, adj_col in adjacent_positions:
                v = final_grid[adj_row, adj_col]
                if v == 0:
                    continue
                if v > 0:
                    final_grid[adj_row, adj_col] = v % 3 + 1
                elif v < 0:
                    final_grid[adj_row, adj_col] = -((-v) % 3 + 1)
        return final_grid


def print_grid(grid: np.ndarray, title: str = ""):
    """打印网格状态（支持任意大小和负数状态）"""
    if title:
        print(f"\n{title}:")

    rows, cols = grid.shape

    # 打印顶部边框
    print("┌" + "─────┬" * (cols - 1) + "─────┐")

    for i in range(rows):
        print("│", end=" ")
        for j in range(cols):
            if grid[i, j] == 0:
                print("   ", end=" │ ")
            else:
                # 显示负数状态
                print(f"{grid[i, j]:2d} ", end=" │ ")
        print()

        # 打印分隔线（除了最后一行）
        if i < rows - 1:
            print("├" + "─────┼" * (cols - 1) + "─────┤")

    # 打印底部边框
    print("└" + "─────┴" * (cols - 1) + "─────┘")


if __name__ == "__main__":
    """主函数 - 演示程序使用"""
    print("=" * 30)
    print("状态说明:")
    print("  0: 无开关")
    print("  1,2,3: 可手动操作的开关（1=关闭, 2=等待, 3=开启）")
    print(" -1,-2,-3: 不可手动操作的开关（-1=关闭, -2=等待, -3=开启）")
    print()

    # 初始状态 - 包含不可手动操作的开关
    grid = [[2, 3, 0], [2, -1, 2], [0, 1, 1]]

    print_grid(np.array(grid), "初始状态")

    solver = SwitchPuzzleSolver(grid)
    moves = solver.solve()

    if moves:
        print(f"\n解决方案：需要按动 {len(moves)} 次")
        # 新增：每一步后输出状态
        current_grid = np.array(grid)
        for i, (row, col) in enumerate(moves, 1):
            print(f"=========第{i}步：按动位置 ({row+1}, {col+1})=========")
            # 模拟这一步
            if solver.is_manual_switch(row, col):
                # 改变当前位置：正负分别循环加一
                v = current_grid[row, col]
                if v > 0:
                    current_grid[row, col] = v % 3 + 1
                elif v < 0:
                    current_grid[row, col] = -((-v) % 3 + 1)

                adjacent_positions = solver.get_adjacent_positions(row, col)
                for adj_row, adj_col in adjacent_positions:
                    v = current_grid[adj_row, adj_col]
                    if v == 0:
                        continue
                    if v > 0:
                        current_grid[adj_row, adj_col] = v % 3 + 1
                    elif v < 0:
                        current_grid[adj_row, adj_col] = -((-v) % 3 + 1)
            print_grid(current_grid, f"第{i}步后的状态")
        final_grid = solver.simulate_moves(moves)
        print_grid(final_grid, "最终状态")

        # 检查是否所有开关都开启
        all_on = True
        for i in range(solver.rows):
            for j in range(solver.cols):
                if final_grid[i, j] != 0 and abs(final_grid[i, j]) != 3:
                    all_on = False
                    break

        if all_on:
            print("\n🎉 成功！所有开关都已开启！")
        else:
            print("\n⚠️  警告：最终状态不正确，可能存在计算错误")
            print("最终状态检查：")
            for i in range(solver.rows):
                for j in range(solver.cols):
                    if final_grid[i, j] != 0:
                        print(f"位置({i+1},{j+1}): {final_grid[i, j]}")
    else:
        print("\n无解！")

    print("\n" + "=" * 50)
